import sys
parent_dir = "../pygl"
sys.path.append(parent_dir)
import time
import OpenGL.GL as gl
import glm
import imgui
import numpy as np
import pygl
from common import AppWithTrackball, AxesHelper, GridHelper

class PointCloudsApp(AppWithTrackball):
    def setupVertex(self) -> None:
        gl.glEnable(gl.GL_DEPTH_TEST)
        self.background_color = glm.vec4(0.3, 0.3, 0.3, 1.0)
        self.program = pygl.ProgramVF("usecase/shaders/point_clouds.vs", 
                                      "usecase/shaders/point_clouds.fs")
        self.max_points = 1000000
        self.readPointCloudsData("dataset/point_clouds.txt")
        self.vertex_vbo = pygl.VertexBufferObject(self.vertices)
        self.vao = pygl.VertexArrayObject()
        binding_point = 0
        self.vao.setVertexBuffer(self.vertex_vbo, binding_point, 0, 
                                 6 * gl.sizeof(gl.GLfloat))
        attribute_position = pygl.VertexAttribute("position", 0, 3, 
                                                  gl.GL_FLOAT, False, 0)
        attribute_color = pygl.VertexAttribute("color",1, 3, gl.GL_FLOAT, 
                                False, 3 * gl.sizeof(gl.GLfloat))
        self.vao.setVertexAttribute(binding_point, attribute_position)
        self.vao.setVertexAttribute(binding_point, attribute_color)
        self.model = glm.mat4(1.0)
        self.camera.far = 800
        self.camera.zoom_speed = 6.0
        self.camera.reset(glm.vec3(2.0, 1.0, 10.0), glm.vec3(0.0, 0.0, 0.0), 
                    glm.vec3(0.0, 1.0, 0.0))
        self.world_axes = AxesHelper(150.0, 5.0)
        self.grid = GridHelper(300.0, 30, glm.vec3(1.0, 0.0, 0.0), 
                               glm.vec3(0.0, 0.0, 1.0))
        self.show_world_axes = True
        self.show_grid = True
        self.imgui_window_width = 250
        self.imgui_window_height = 260

    # 配置imgui界面元素
    def showImGui(self)->None:
        imgui.begin("设置")
        imgui.text("文件读取耗时：{:.2f}秒".format(self.read_end_time - self.start_time))
        imgui.text("数据转换耗时：{:.2f}秒".format(self.process_end_time - 
                                           self.read_end_time))
        if imgui.button("重置"):
            self.camera.reset(glm.vec3(2.0, 1.0, 10.0), glm.vec3(0.0, 0.0, 0.0), 
                    glm.vec3(0.0, 1.0, 0.0))
        imgui.text("轨迹球位置")
        imgui.begin_table("position data", 4)
        imgui.table_next_row()
        for j in range(3):
            imgui.table_next_column()
            imgui.text("{:.2f}".format(self.camera.position[j]))
        imgui.end_table()
        _, self.max_points = imgui.slider_int("观察点数", self.max_points, 
                                              1000000, self.total_points)
        _, self.camera.zoom_speed = imgui.slider_float("缩放速度", 
                                    self.camera.zoom_speed, 4.0, 10.0)
        _, self.show_world_axes = imgui.checkbox("显示三维世界坐标轴", 
                                                self.show_world_axes)
        _, self.show_grid = imgui.checkbox("显示XZ平面网格", 
                                                self.show_grid)
        imgui.end()   

    def render(self) -> None:
        gl.glClearNamedFramebufferfv(0, gl.GL_COLOR, 0, 
                                     self.background_color.to_tuple())
        gl.glClearNamedFramebufferfv(0, gl.GL_DEPTH, 0, [1.0]); 
        self.program.use()
        view = self.camera.getViewMatrix()
        projection = self.camera.getProjectionMatrix()
        self.program.setUniformMatrix4fv(0, self.model)
        self.program.setUniformMatrix4fv(1, view)
        self.program.setUniformMatrix4fv(2, projection)
        self.vao.bind()
        gl.glPointSize(2.0)
        gl.glDrawArrays(gl.GL_POINTS, 0, self.max_points)
        if self.show_world_axes:
            self.world_axes.updateModelMatrix
            self.world_axes.render(view, projection)
        if self.show_grid:
            self.grid.render(view, projection)

    def cleanup(self) -> None:
        self.vao.delete()
        self.vertex_vbo.delete()
        self.program.delete()

    def readPointCloudsData(self, file_name:str)->None:
        self.start_time = time.time()
        with open(file_name) as file:
            lines = file.readlines()
            self.total_points = len(lines)
            if self.max_points > self.total_points:
                self.max_points = self.total_points
            self.read_end_time = time.time()
            # 创建点云三维顶点数组
            self.vertices = np.empty((self.total_points, 6), dtype=gl.GLfloat)
            for index in range(self.total_points):
                data = lines[index].split()
                # 读取三维坐标 x y z
                self.vertices[index, 0] = float(data[1])
                self.vertices[index, 1] = float(data[2])
                self.vertices[index, 2] = float(data[0])
                # 读取颜色数据 R G B        
                self.vertices[index, 3] = float(data[4]) / 255.0
                self.vertices[index, 4] = float(data[5]) / 255.0
                self.vertices[index, 5] = float(data[6]) / 255.0
        self.process_end_time = time.time()

if __name__ == "__main__":
    app = PointCloudsApp(100, 100, 1000, 600, "你好，点云App！")
    app.run()


